"""PlotStack templates."""
import plots
from block import drawers
from analysis import tdind
import pyqtgraph as gr
import PyQt4.QtCore as QtCore
# ATTENTION! Do not forget to add a data block!!!


class MainTemplate(plots.PlotStack):
    """Some features to be used by descendant classes, such as pre-defined
    collections of moving averages."""
    def _AddEMA(self, plot, dataSource, size, color, fieldName="close"):
        """Adde EMA to given plot and dataSource."""
        ema = tdind.EMA()
        ema["input"].fieldName = fieldName
        ema.size = size
        pd = drawers.PathDrawer(plot)
        pd.curve.setPen(color=color, width=1)
        dataSource.Connect(ema)
        ema.Connect(pd)
        
    def Add8To1EMAs(self, **kwargs):
        """Adds 8-to-1 moving averages.
        
        Arguments:
            plot= -- (Optional) NavalaPlot, defaults to self.plot1
            dataSource -- (Optional) PlotDataBlock, defaults to
                self.plot1._dataBlocks["dataSource"].
            shortestSize -- (Optional) shortest EMA size. Defaults to 3
        """
        plot = kwargs.get("plot", self.plot1)
        dataSource = kwargs.get("dataSource", plot._dataBlocks["dataSource"])
        shortestSize = kwargs.get("shortestSize", 3)
        longestSize = shortestSize*8
        
        self._AddEMA(plot=plot, dataSource=dataSource, size=shortestSize, color=(216, 216, 216))
        self._AddEMA(plot=plot, dataSource=dataSource, size=longestSize, color=(160, 160, 160))


    

    def Add7EMAs(self, **kwargs):
        """Adds 7 moving averages.
        
        Arguments:
            plot= -- (Optional) NavalaPlot, defaults to self.plot1
            dataSource -- (Optional) PlotDataBlock, defaults to
                self.plot1._dataBlocks["dataSource"].
        """
        plot = kwargs.get("plot", self.plot1)
        dataSource = kwargs.get("dataSource", plot._dataBlocks["dataSource"])
        
        n, inc = 7, 32
        yy = range(255-(n-1)*inc, 256, inc)[-1::-1] # For the colors
        sizes = [3, 5, 8, 13, 21, 34, 55]
        for (i, (y, size)) in enumerate(zip(yy, sizes)):
            self._AddEMA(plot=plot, dataSource=dataSource, 
                size=size, color=(0, y, y))
            
    def AddMAC(self, **kwargs):
        """Adds J Bernstein's Moving Average Channel.
        
        Arguments:
            plot=self.plot1 -- (Optional) NavalaPlot, defaults to self.plot1
            dataSource -- (Optional) PlotDataBlock, defaults to
                self.plot1._dataBlocks["dataSource"].
            size_high=10 -- size of Moving Average (MA) of highs
            size_low=8 -- size of MA of lows
        """
        plot = kwargs.get("plot", self.plot1)
        dataSource = kwargs.get("dataSource", plot._dataBlocks["dataSource"])
        size_high = kwargs.get("size_high", 10)
        size_low = kwargs.get("size_low", 8)
        
        colors = ((124, 192, 207), (121, 189, 204))
        sizes = [size_high, size_low]
        fieldNames = ["high", "low"]
        for (color_, size, fieldName) in zip(colors, sizes, fieldNames):
            self._AddEMA(plot=plot, dataSource=dataSource, 
                size=size, color=color_, fieldName=fieldName)
                
    def Add2EMAs(self, **kwargs):
        """Adds 2 EMAs
        
        Arguments:
            plot=self.plot1 -- (Optional) NavalaPlot, defaults to self.plot1
            dataSource -- (Optional) PlotDataBlock, defaults to
                self.plot1._dataBlocks["dataSource"].
            size1=17 -- size of short Moving Average (MA)
            size_low=72 -- size of long MA
        
        Reference: Aliakyn seminar
        """
        plot = kwargs.get("plot", self.plot1)
        dataSource = kwargs.get("dataSource", plot._dataBlocks["dataSource"])
        size1 = kwargs.get("size1", 17)
        size2 = kwargs.get("size2", 72)
        
        colors = ((184, 42, 242), (242, 239, 42))
        sizes = [size1, size2]
        fieldNames = ["close", "close"]
        for (color_, size, fieldName) in zip(colors, sizes, fieldNames):
            self._AddEMA(plot=plot, dataSource=dataSource, 
                size=size, color=color_, fieldName=fieldName)
                
    def AddHiLoPaths(self, **kwargs):
        """Adds high and low paths.
        
        Arguments:
            plot=self.plot1 -- (Optional) NavalaPlot, defaults to self.plot1
            dataSource -- (Optional) PlotDataBlock, defaults to
                self.plot1._dataBlocks["dataSource"].
        """
        plot = kwargs.get("plot", self.plot1)
        dataSource = kwargs.get("dataSource", plot._dataBlocks["dataSource"])
        
        pdhi = drawers.PathDrawer(plot)
        pdhi["input"].fieldName = "high"
        pdhi.curve.setPen(color=(0, 192, 0), width=1)
        dataSource.Connect(pdhi)
    
        pdlo = drawers.PathDrawer(plot)
        pdlo.curve.setPen(color=(255, 0, 0), width=1)
        pdlo["input"].fieldName = "low"
        dataSource.Connect(pdlo)



class Candle(MainTemplate):
    def __init__(self, *args, **kwargs):
        plots.PlotStack.__init__(self, *args, **kwargs)
        # Candle plot
        p1 = plots.NavalaPlot(self)
#        p1.hideAxis("bottom") # TODO How to get rid of this now
        ds1 = p1.AddDataBlock("dataSource")

#        ema1 = tdind.EMA()
#        ema1["input"].fieldName = "close"
#        ema1.size = 17
#        ema2 = tdind.EMA()
#        ema2["input"].fieldName = "close"
#        ema2.size = 34

        cd = drawers.CandlestickDrawer(p1)
#        pd1 = drawers.PathDrawer(p1)
#        pd1.curve.setPen(color=(0, 192, 192), width=1)
#        pd2 = drawers.PathDrawer(p1)
#        pd2.curve.setPen(color=(128, 128, 64), width=2)

        ds1.Connect(cd)
#        ds1.Connect(ema1)
#        ds1.Connect(ema2)
#        ema1.Connect(pd1)
#        ema2.Connect(pd2)
        
        self.nextRow()
        # This is to define the sharing of space between the price and volume
        # plot.
        self.layout.setRowStretchFactor(1, 210)
        
###        for i in [1, 2]:
###            self.layout.setRowPreferredHeight(i, 0)
###            self.layout.setRowMinimumHeight(i, 0)
###            self.layout.setRowSpacing(i, 0)
######        self.layout.setRowMinimumHeight(2, 1)
###        self.layout.setRowPreferredHeight(2, 20)
####        self.layout.setRowMaximumHeight(2, 20)
###        p2.layout.setRowMaximumHeight(2, 20) # Tries to set maximum height of viewbox inside PlotItem #StretchFactor(2, 100)

###        self.layout.setRowPreferredHeight(1, 10)
###        self.layout.setRowPreferredHeight(2, 10)
        
        # Price plot
        self.plot1 = p1
        # Data source blocks
        self.dataSource1 = ds1
        
class Candle2(Candle):
    """Candle Moving Average Channel"""
    def __init__(self, *args, **kwargs):
        Candle.__init__(self, *args, **kwargs)
#        self.AddMAC(size_high=10, size_low=8)
#        self._AddEMA(plot=self.plot1, dataSource=self.dataSource1, 
#            size=34, color=(78, 147, 96))
#        self.AddHiLoPaths()
        self.Add2EMAs()
        
######        pp = tdind.LSPredictor()
######        pp["input"].fieldName = "close"
######        pp.trainingSize = 30
######        pp.order = 6
######        pp.travel = 5
######
######        pd = drawers.PathDrawer(self.plot1)
######        pd["input"].fieldName = "y"
######        pd.curve.setPen(color=(216, 216, 216), width=1)
######
######        self.dataSource1.Connect(pp)
######        pp.Connect(pd)
        
    
class CandleMomentum(Candle):
    """Candle plus momentum indicator"""
    def __init__(self, *args, **kwargs):
        Candle.__init__(self, *args, **kwargs)
        
        self.Add2EMAs()

        # Price plot
        self.plot1.hideAxis("bottom") # TODO How to get rid of this now

        # Momentum plot
        p2 = plots.NavalaPlot(self)
        ds2 = p2.AddDataBlock("dataSource")
        self.plot2 = p2
        self.dataSource2 = ds2

        # Momentum block and drawer
        om = tdind.Momentum()
        om["input"].fieldName = "close"
        om.size = 21
        pd = drawers.PathHLinesDrawer(p2, numLines=1, linePositions=[0])
        pd.curve.setPen(color=(255, 255, 255), width=1)
        ds2.Connect(om)
        om.Connect(pd)


        # This is to define the sharing of space between the two plots
        self.layout.setRowStretchFactor(1, 210)
        self.layout.setRowStretchFactor(2, 105)
        
    
        
class CandleVolume(Candle):
    def __init__(self, *args, **kwargs):
        Candle.__init__(self, *args, **kwargs)
        
        # Price plot
        self.plot1.hideAxis("bottom") # TODO How to get rid of this now

        # Volume plot
        p2 = plots.NavalaPlot(self)
        ds2 = p2.AddDataBlock("dataSource")
        vd = drawers.VolumeDrawer(p2)
        ds2.Connect(vd)

        # This is to define the sharing of space between the price and volume
        # plot.
        self.layout.setRowStretchFactor(1, 210)
        self.layout.setRowStretchFactor(2, 30)
        
        for i in [1, 2]:
            self.layout.setRowPreferredHeight(i, 0)
            self.layout.setRowMinimumHeight(i, 0)
            self.layout.setRowSpacing(i, 0)
###        self.layout.setRowMinimumHeight(2, 1)
        self.layout.setRowPreferredHeight(2, 20)
#        self.layout.setRowMaximumHeight(2, 20)
        p2.layout.setRowMaximumHeight(2, 20) # Tries to set maximum height of viewbox inside PlotItem #StretchFactor(2, 100)

###        self.layout.setRowPreferredHeight(1, 10)
###        self.layout.setRowPreferredHeight(2, 10)
        
        # Volume plot
        self.plot2 = p2
        # Data source blocks
        self.dataSource2 = ds2
        
        self.Add2EMAs()
        
class PathVolume(MainTemplate):
    def __init__(self, *args, **kwargs):
        super(self.__class__, self).__init__(*args, **kwargs)
        p1 = plots.NavalaPlot(self)
        p1.hideAxis("bottom") # TODO How to get rid of this now
        self.nextRow()
        p2 = plots.NavalaPlot(self) # Funny but this is it (this is the volume plot)
        self.layout.setRowStretchFactor(1, 21)
        self.layout.setRowStretchFactor(2, 8)

        ds1 = p1.AddDataBlock("dataSource")
        ds2 = p2.AddDataBlock("dataSource")


        # Adds drawers
        pd = drawers.PathDrawer(p1)
        pd["input"].fieldName = "close"
        vd = drawers.VolumeDrawer(p2)
        
        ds1.Connect(pd)
        ds1.Connect(vd)
        
        # Price plot
        self.plot1 = p1
        # Volume plot
        self.plot2 = p2

        # Data source blocks
        self.dataSource1 = ds1
        self.dataSource2 = ds2


class CVMA17(CandleVolume):
    """Candle Volume Moving Average"""
    def __init__(self, *args, **kwargs):
        CandleVolume.__init__(self, *args, **kwargs)
        
        ema1 = tdind.EMA()
        ema1["input"].fieldName = "close"
        ema1.size = 17

        pd1 = drawers.PathDrawer(self.plot1)
        pd1.curve.setPen(color=(0, 192, 192), width=1)

        self.dataSource1.Connect(ema1)

        ema1.Connect(pd1)

class CV7EMAs(CandleVolume):
    """Candle Volume Moving Average"""
    def __init__(self, *args, **kwargs):
        CandleVolume.__init__(self, *args, **kwargs)

        self.Add7EMAs()

class CV8To1EMAs(CandleVolume):
    """Candle Volume Moving Average"""
    def __init__(self, *args, **kwargs):
        CandleVolume.__init__(self, *args, **kwargs)

        self.Add8To1EMAs(shortestSize=3)

class CVMAC(CandleVolume):
    """Candle Volume Moving Average"""
    def __init__(self, *args, **kwargs):
        CandleVolume.__init__(self, *args, **kwargs)
        self.AddMAC(size_high=10, size_low=8)
    

class CVStochastic(CandleVolume):
    """Adds a third Stochastic plot.
    
    The number of days is 14. The Stochastic indicator (SI) will be the "slow"
    one (3-EMA of original SI). A second dashed line will be the 3-EMA of the
    slow SI. This is set according to the SI system described in [1].
    
    References:
    [1] Jake Bernstein - Complete Day Trader Vol 1
    """
    def __init__(self, *args, **kwargs):
        CandleVolume.__init__(self, *args, **kwargs)

#        self.Add8To1EMAs(shortestSize=3)
#        self.AddMAC(size_high=10, size_low=8)
        self.AddHiLoPaths()
        
        self.plot2.hideAxis("bottom") # TODO How to get rid of this now
        self.nextRow()
        # Volume plot
        p3 = plots.NavalaPlot(self)
        ds3 = p3.AddDataBlock("dataSource")
        # This is to define the share of space that the stochastic plot is going to occupy
        self.layout.setRowStretchFactor(3, 5)
        
#        sto1 = tdind.Stochastic()
#        sto1.size = 5 # Check this later, don't know exactly which size to apply
#
#        sto2 = tdind.Stochastic()
#        sto2.size = 9 # Check this later, don't know exactly which size to apply

        sto3 = tdind.Stochastic()
        sto3.size = 14 # Jake Bernstein - Complete DT Vol 1

        ema1 = tdind.EMA()
#        ema1["input"].fieldName = "y"
        ema1.size = 3

        ema2 = tdind.EMA()
#        ema2["input"].fieldName = "close"
        ema2.size = 3
        
#        pd1 = drawers.PathDrawer(p3)
#        pd1.curve.setPen(color=(255, 128, 0), width=1)
#        ds3.Connect(sto1)
#        sto1.Connect(pd1)
#
#        pd2 = drawers.PathDrawer(p3)
#        pd2.curve.setPen(color=(255, 196, 0), width=1)
#        ds3.Connect(sto2)
#        sto2.Connect(pd2)

        pd31 = drawers.PathDrawer(p3)
        pd31.curve.setPen(color=(255, 255, 0), width=1)
        
        pd32 = drawers.PathHLinesDrawer(p3, numLines=2, linePositions=[.25, .75])
        pd32.curve.setPen(color=(255, 255, 0), width=1, style=QtCore.Qt.DashLine)

        # Makes the circuit
        ds3.Connect(sto3)
        sto3.Connect(ema1)
        ema1.Connect(pd31)
        ema1.Connect(ema2)
        ema2.Connect(pd32)

        # Stochastic plot
        self.plot3 = p3
        # Data source block
        self.dataSource3 = ds3


        for i in [1, 2, 3]:
            self.layout.setRowPreferredHeight(i, 1)
            self.layout.setRowMinimumHeight(i, 0)
#            self.layout.setRowMaximumHeight(i, 0)
            self.layout.setRowSpacing(i, 0)
